import math

import numpy as np
import torch

from algorithms.algorithm import Algorithm
from algorithms.base.nsga2 import NSGA2
from indicators.hv import HV
from individual.population import Population
from operators.crossover.de import DE
from operators.crossover.sbx import SBX
from operators.mutations.polynomial_mutation import PolynomialMutation
from problems.problem import Problem
from utils.select_operator import tournament_selection


class LSMOF(Algorithm):
    parameters = [
        {'label': 'wD', 'name': 'wd', 'type': 'number', 'description': 'The generation of weight optimization with DE',
         'step': 1, 'defaultValue': 10},
        {'label': 'SubN', 'name': 'sub_n', 'type': 'number', 'step': 1,
         'description': 'The population size of the transferred problem', 'defaultValue': 30},
        {
            'label': 'operator', 'name': 'operator', 'type': 'select', 'description': 'Original operators',
            'defaultValue': 0,
            "options": [
                {'label': "GA", "value": 0},
                {'label': "DE", "value": 1}
            ]
        }
    ]

    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs,
                 ):
        super().__init__(pop_size, problem, kwargs)
        self.population: Population = problem.init_population(pop_size)
        self.g = math.ceil(problem.max_fun_eval * 0.05 / (self.sub_n * 2 * self.wd))
        self.archive: Population = None

    def each_iteration(self):
        if self.problem.fun_eval < 0.6 * self.problem.max_fun_eval:
            self.archive = self.weight_optimization(self.problem, self.g, self.population, self.wd, self.sub_n)
            temp_population = self.population + self.archive
            self.population = self.environmental_selection(temp_population, self.pop_size)[0]
        else:
            self.population = self.sub_NSGA2(self.problem, self.population, self.operator, self.pop_size)

    @classmethod
    def weight_optimization(cls, problem: Problem, g, population: Population, wd, sub_n):
        obj = population.obj
        reference = torch.max(obj, 0)[0].unsqueeze(0)
        ref_pop, _, _ = cls.environmental_selection(population, wd)

        # 计算参考方向
        direction = torch.cat((torch.sqrt(torch.sum((ref_pop.dec - problem.low_limit.repeat(wd, 1)) ** 2, dim=1)),
                               torch.sqrt(
                                   torch.sum((problem.high_limit.repeat(wd, 1) - ref_pop.dec) ** 2, dim=1))))

        direct = torch.cat(((ref_pop.dec - problem.low_limit.repeat(wd, 1)),
                            (problem.high_limit.repeat(wd, 1) - ref_pop.dec))) / direction.unsqueeze(1).repeat(1,
                                                                                                               problem.var_dim)
        w_max = torch.sqrt(torch.sum((problem.high_limit - problem.low_limit) ** 2)) * 0.5

        # 初始化种群
        w0 = torch.rand(sub_n, 2 * wd) * w_max
        fitness, pop_new = cls.function_fitness(problem, w0, direct, reference)
        archive = pop_new[torch.where(NSGA2.ND_sort(pop_new.obj, 1)[0] == 1)[0]]

        pCR = 0.2
        beta_min, beta_max = 0.2, 0.8

        for it in range(g):
            for i in range(sub_n):
                x = w0[i, :]
                A = np.delete(np.arange(sub_n), i)
                np.random.shuffle(A)
                a, b, c = A[:3]

                # 变异
                beta = (beta_max - beta_min) * torch.rand(2 * wd) + beta_min
                y = w0[a, :] + beta * (w0[b, :] - w0[c, :])
                y = torch.clamp(y, 0, w_max)

                # 交叉
                z = x.clone()
                j0 = torch.randint(0, x.numel(), (1,))
                for j in range(x.numel()):
                    if j == j0 or torch.rand(1) <= pCR:
                        z[j] = y[j]
                    else:
                        z[j] = x[j]

                # 评估新解
                fit, _ = cls.function_fitness(problem, z.unsqueeze(0), direct, reference)
                if fit < fitness[i]:
                    w0[i, :] = z
                    fitness[i] = fit
        # 更新并存储非支配解
        return archive[torch.where(NSGA2.ND_sort(archive.obj, 1)[0] == 1)[0]]

    @staticmethod
    def environmental_selection(population, pop_size):
        """
        This function performs the environmental selection step for LSMOF.
        """
        # Non-dominated sorting
        front_no, MaxFNo = NSGA2.ND_sort(population.obj, pop_size)
        next_index = torch.zeros_like(front_no, dtype=torch.bool)
        next_index[front_no < MaxFNo] = True

        # Calculate the crowding distance of each solution
        crowd_dis = NSGA2.crowding_distance(population.obj, front_no)

        # Select the solutions in the last front based on their crowding distances
        Last = torch.where(front_no == MaxFNo)[0]
        _, Rank = torch.sort(crowd_dis[Last], descending=True)
        selection_count = pop_size - next_index.sum().item()
        next_index[Last[Rank[:selection_count]]] = True

        # population for next_index generation
        front_no = front_no[next_index]
        crowd_dis = crowd_dis[next_index]
        population = population[next_index]

        return population, front_no, crowd_dis

    @classmethod
    def sub_NSGA2(cls, problem, population, operator, pop_size):
        """
        Sub-optimizer in LSMOF (NSGA-II)
        """
        obj = population.obj
        dec = population.dec
        # Non-dominated sorting
        FrontNo, _ = NSGA2.ND_sort(obj, float('inf'))

        # Calculating crowding distance
        CrowdDis = NSGA2.crowding_distance(obj, FrontNo)

        # Selection using tournament based on FrontNo and CrowDis
        # Apply genetic operators based on the specified operator
        if operator == 1:
            MatingPool = tournament_selection(2, pop_size, [FrontNo, -CrowdDis])
            Offspring = cls.operator_GA(problem, dec[MatingPool])
        else:
            MatingPool1 = tournament_selection(2, pop_size, [FrontNo, -CrowdDis])
            MatingPool2 = tournament_selection(2, pop_size, [FrontNo, -CrowdDis])
            MatingPool3 = tournament_selection(2, pop_size, [FrontNo, -CrowdDis])
            Offspring = cls.DE(problem, dec[MatingPool1], dec[MatingPool2], dec[MatingPool3])

        # Combine the parent and offspring populations, and perform environmental selection
        CombinedPopulation = population + Offspring
        NextPopulation = cls.environmental_selection(CombinedPopulation, pop_size)[0]
        return NextPopulation

    @staticmethod
    def function_fitness(problem: Problem, w0, direct, reference):
        # 获取 w0 的形状信息
        sub_n, wd = w0.shape
        wd = wd // 2

        # 初始化 obj
        obj = torch.zeros(sub_n, 1)

        pop_dec_list = []
        pop_obj_list = []

        for i in range(sub_n):
            pop_dec = torch.cat((
                w0[i, :wd].unsqueeze(1).repeat(1, problem.var_dim) * direct[:wd, :] + problem.low_limit.repeat(wd, 1),
                problem.high_limit.repeat(wd, 1) - w0[i, wd:].unsqueeze(1).repeat(1, problem.var_dim) * direct[wd:, :]
            ))
            pop_dec_list.append(pop_dec)

            off_w_pop = problem.estimate_population(pop_dec).obj
            pop_obj_list.append(off_w_pop)

            obj[i] = -HV.value(off_w_pop, reference)
        pop_dec = torch.cat(pop_dec_list, dim=0)
        pop_obj = torch.cat(pop_obj_list, dim=0)
        return obj, Population(pop_dec, pop_obj)

    @classmethod
    def operator_GA(cls, problem, dec):
        dec = SBX.do(dec, dec, problem, offspring_size=1, pro_c=1, dis_c=20)
        dec = PolynomialMutation.do(dec, problem)
        obj = problem.estimate_population(dec).obj
        return Population(dec, obj)

    @classmethod
    def DE(cls, problem, dec1, dec2, dec3):
        dec = DE.do(problem, dec1, dec2, dec3)
        obj = problem.estimate_population(dec).obj
        return Population(dec, obj)
